import time
import threading
from io import BytesIO
import pyautogui
import mss
from PIL import Image, ImageDraw
from flask import Flask, Response
from flask_sockets import Sockets
from flask_cors import CORS
import queue

app = Flask(__name__)
sockets = Sockets(app)
CORS(app, supports_credentials=True)

# 优化配置
CONFIG = {
  'monitor': 2,
  'quality': 90,  # 进一步降低质量
  'scale_factor': 0.8,  # 进一步缩小
  'target_fps': 30,
  'buffer_size': 2  # 减少缓冲大小以降低延迟
}

# 全局变量
latest_frame = None
frame_lock = threading.Lock()
running = True
clients = set()

arrow_img = Image.open("aero_arrow_xl.png")
arrow_img = arrow_img.convert("RGBA")


def get_monitor_config():
  with mss.mss() as sct:
    if CONFIG['monitor'] >= len(sct.monitors):
      CONFIG['monitor'] = 1

    mon = sct.monitors[CONFIG['monitor']]
    return {
      "top": mon["top"],
      "left": mon["left"],
      "width": mon["width"],
      "height": mon["height"],
      "mon": CONFIG['monitor'],
    }


def capture_frames():
  """优化的截屏线程"""
  global latest_frame, running

  monitor_config = get_monitor_config()
  original_width = monitor_config['width']
  original_height = monitor_config['height']

  scaled_width = int(original_width * CONFIG['scale_factor'])
  scaled_height = int(original_height * CONFIG['scale_factor'])

  frame_time = 1.0 / CONFIG['target_fps']

  with mss.mss() as sct:
    while running:
      try:
        start_time = time.time()

        # 截屏
        sct_img = sct.grab(monitor_config)

        # 转换为PIL并缩放
        img = Image.frombytes('RGB', sct_img.size, sct_img.rgb)

        if CONFIG['scale_factor'] != 1.0:
          img = img.resize((scaled_width, scaled_height), Image.Resampling.LANCZOS)

        # 绘制鼠标
        x, y = pyautogui.position()
        rel_x = x - monitor_config['left']
        rel_y = y - monitor_config['top']

        if (0 <= rel_x < original_width and
            0 <= rel_y < original_height):
          scaled_x = int(rel_x * CONFIG['scale_factor'])
          scaled_y = int(rel_y * CONFIG['scale_factor'])

          img.paste(arrow_img, (scaled_x, scaled_y), arrow_img)

        # 使用更快的压缩参数
        buffer = BytesIO()
        img.save(buffer, 'JPEG',
                 quality=CONFIG['quality'],
                 optimize=True,
                 progressive=False,  # 非渐进式，解码更快
                 subsampling=2)  # 更高的色度子采样
        frame_data = buffer.getvalue()

        # 更新帧
        with frame_lock:
          latest_frame = frame_data

        # 精确的帧率控制
        elapsed = time.time() - start_time
        sleep_time = max(0, frame_time - elapsed)
        if sleep_time > 0:
          time.sleep(sleep_time)

      except Exception as e:
        print(f"截屏错误: {e}")
        time.sleep(0.01)


@app.route('/')
def index():
  with open('index.html', encoding='utf-8') as f:
    return f.read()


@sockets.route('/display')
def display(ws):
  print('WebSocket 连接建立')
  clients.add(ws)

  try:
    # 发送第一帧
    with frame_lock:
      if latest_frame is not None:
        ws.send(latest_frame)

    while not ws.closed and running:
      # 等待客户端请求
      msg = ws.receive()
      if msg == 'next':
        with frame_lock:
          if latest_frame is not None:
            ws.send(latest_frame)
      elif msg is None:
        break

  except Exception as e:
    print(f"WebSocket 错误: {e}")
  finally:
    clients.discard(ws)
    print('WebSocket 连接关闭')


if __name__ == '__main__':
  capture_thread = threading.Thread(target=capture_frames, daemon=True)
  capture_thread.start()

  try:
    from gevent import pywsgi
    from geventwebsocket.handler import WebSocketHandler

    server = pywsgi.WSGIServer(
      ('0.0.0.0', 80),
      app,
      handler_class=WebSocketHandler
    )
    print('服务器启动: http://0.0.0.0:80')
    server.serve_forever()
  except KeyboardInterrupt:
    print('正在关闭服务器...')
    running = False